import cv2 as cv
import mxnet as mx
import numpy as np
import os.path as osp
from math import ceil
from nn import get_net
import utils

__all__ = ['MtcnnV2', 'load_model']

def load_model(name, model, ctx=mx.cpu()):
    if not model:
        return None
    net = get_net(name)
    net.load_parameters(model)
    net.hybridize()
    data = mx.sym.var('data', shape=(None, 3, net.size, net.size))
    sym = mx.sym.Group(net(data))
    mod = mx.mod.Module(symbol=sym, context=ctx, label_names=None)
    mod.bind(data_shapes=[('data', (1, 3, net.size, net.size))], for_training=False)  
    arg_params = {}
    aux_params = {}
    for key, param in net.collect_params().items():
        if key.endswith('running_mean') or key.endswith('running_var'):
            aux_params[key] = param.data().copy()
        else:
            arg_params[key] = param.data().copy()
    mod.set_params(arg_params, aux_params, allow_missing=False)
    return mod

class MtcnnV2(object):
    """MTCNN face detector.
    
    Parameters
    ----------
    min_face : int, default is 20
        Minimum size of face to be detected.
    scale_factor : float, default is 0.709
        Scale factor for image pyramid in Pnet.
    thresholds : list of float, default is [0.6, 0.7, 0.7]
        Thresholds for Pnet, Rnet and Onet.
    """
    def __init__(self, pnet=None, rnet=None, onet=None,
        min_face=20, scale_factor=0.709, thresholds=[0.6, 0.7, 0.7]):
        self._min_face = min_face
        self._scale_factor = scale_factor
        self._thresholds = thresholds
        # load models
        # self.pnet = utils.load_model('models/caffe/pnet', 0, 12, ctx)
        # self.pnet = self._load_model('pnet', pnet, ctx)
        # self.rnet = self._load_model('rnet', rnet, ctx)
        # self.onet = self._load_model('onet', onet, ctx)
        self.pnet = pnet
        self.rnet = rnet
        self.onet = onet
    
    def detect(self, image):
        """Detect faces with Pnet, Rnet and Onet.
        
        Parameters
        ----------
        image : numpy.ndarray, h x w x c [RGB]
        
        Example:
        ```
        mtcnn = Mtcnn("models/caffe", min_face=40)
        image = cv.imread('data/test.jpg', cv.IMREAD_COLOR)
        image = cv.cvtColor(image, cv.COLOR_BGR2RGB)
        bbox = mtcnn.detect(image)
        utils.plot_bbox(image, bbox)
        ```

        Returns:
        -------
            numpy.ndarray, n x 5 [score, x1, y1, x2, y2]
        """
        im = (image.astype(np.float32) - 127.5) * 0.0078125 # normalize

        ## Pnet
        if not self.pnet:
            return np.array([])
        min_len = min(im.shape[:2])
        scales = utils.scale_pyramid(min_len, self._min_face, self._scale_factor)
        total_bbox = []
        for scale in scales:
            width = int(ceil(im.shape[1] * scale))
            height = int(ceil(im.shape[0] * scale))
            img = cv.resize(im, (width, height))
            img = img.transpose((2, 0, 1))[None, ...] # HWC -> 1CHW
            self.pnet.forward(mx.io.DataBatch([mx.nd.array(img)]))
            cls_pred, bbox_pred = self.pnet.get_outputs()
            bbox = utils.generate_bbox(cls_pred, bbox_pred, scale, self._thresholds[0])
            if bbox.size > 0:
                bbox = bbox[utils.bbox_nms(bbox, 0.5, 'union')]
                total_bbox.append(bbox)
        if len(total_bbox) == 0:
            return np.array([])
        total_bbox = np.vstack(total_bbox)
        total_bbox = total_bbox[utils.bbox_nms(total_bbox, 0.7, 'union')]
        total_bbox = utils.bbox_refine(total_bbox)

        ## Rnet
        if not self.rnet:
            return total_bbox
        data = utils.crop_pad(im, total_bbox, 24)
        data_iter = mx.io.NDArrayIter(data, batch_size=min(len(data), 100))
        cls_pred = []
        bbox_pred = []
        for batch in data_iter:
            self.rnet.forward(batch)
            cls_pred.append(self.rnet.get_outputs()[1].asnumpy())
            bbox_pred.append(self.rnet.get_outputs()[0].asnumpy())
        num = len(total_bbox)
        cls_pred = np.vstack(cls_pred)[:num]
        bbox_pred = np.vstack(bbox_pred)[:num]
        total_bbox[:, 0] = cls_pred[:, 1]
        total_bbox = np.hstack((total_bbox, bbox_pred))
        keep = np.where(total_bbox[:,0] >= self._thresholds[1])[0]
        if len(keep) == 0:
            return np.array([])
        total_bbox = total_bbox[keep]
        total_bbox = total_bbox[utils.bbox_nms(total_bbox, 0.7, 'union')]
        total_bbox = utils.bbox_refine(total_bbox)
        
        ## onet
        if not self.onet:
            return total_bbox
        data = utils.crop_pad(im, total_bbox, 48)
        data_iter = mx.io.NDArrayIter(data, batch_size=min(len(data), 20))
        cls_pred = []
        bbox_pred = []
        for batch in data_iter:
            self.onet.forward(batch)
            cls_pred.append(self.onet.get_outputs()[1].asnumpy())
            bbox_pred.append(self.onet.get_outputs()[0].asnumpy())
        num = len(total_bbox)
        cls_pred = np.vstack(cls_pred)[:num]
        bbox_pred = np.vstack(bbox_pred)[:num]
        total_bbox[:, 0] = cls_pred[:, 1]
        total_bbox = np.hstack((total_bbox, bbox_pred))
        keep = np.where(total_bbox[:,0] >= self._thresholds[1])[0]
        if len(keep) == 0:
            return np.array([])
        total_bbox = total_bbox[keep]
        total_bbox = utils.bbox_refine(total_bbox)
        total_bbox = total_bbox[utils.bbox_nms(total_bbox, 0.7, 'min')]
        return total_bbox

if __name__ == '__main__':
    mtcnn = MtcnnV2(pnet=load_model('pnet', 'models/mtcnn/pnet/032.params'))
    image = cv.imread('data/test.jpg', cv.IMREAD_COLOR)
    image = cv.cvtColor(image, cv.COLOR_BGR2RGB)
    bbox = mtcnn.detect(image)
    utils.plot_bbox(image, bbox)